﻿#include <iostream>
#include <thread>

#include <zmq.hpp>

#include "../../inc/helper.hpp"

const static int g_client_number = 10;
const static int g_worker_number = 3;
const static char* g_fired_msg = "Fired!";

void client(zmq::context_t& ctx)
{
	zmq::socket_t client(ctx, zmq::socket_type::req);
	client.connect("inproc://frontend");

	client.send("HELLO", 5);
	std::string reply = receive_string(client);
	info("Client : {}", reply);
}

void worker(zmq::context_t& ctx)
{
	zmq::socket_t worker(ctx, zmq::socket_type::req);
	worker.connect("inproc://backend");

	worker.send("READY", 5);

	std::random_device rd;
	std::uniform_int_distribution<int> dist(1, 500);

	int total = 0;
	while (true)
	{
		auto requests = zmsg_recv(worker);
		std::string req_str = to_string(requests.back());
		info("Worker: {}", req_str);

		std::vector<zmq::message_t> replies;
		replies.emplace_back(std::move(requests[0]));
		replies.emplace_back();
		replies.emplace_back(to_message("OK"));
		zmsg_send(worker, replies.begin(), replies.end());
	}
}

using thread_ptr = std::shared_ptr<std::thread>;

int main()
{
	zmq::context_t ctx;
	zmq::socket_t frontend(ctx, zmq::socket_type::router);
	zmq::socket_t backend(ctx, zmq::socket_type::router);
	frontend.bind("inproc://frontend");
	backend.bind("inproc://backend");

	std::vector<thread_ptr> threads;
	for (unsigned i = 0; i < g_worker_number; ++i)
	{
		threads.push_back(std::make_shared<std::thread>(worker, std::ref(ctx)));
	}
	for (unsigned i = 0; i < g_client_number; ++i)
	{
		threads.push_back(std::make_shared<std::thread>(client, std::ref(ctx)));
	}

	std::vector<zmq::pollitem_t> items = {
		{backend, 0, ZMQ_POLLIN, 0},
		{frontend, 0, ZMQ_POLLIN, 0}
	};
	std::list<zmq::message_t> worker_queue;
	while (true)
	{
		/// 没有工作线程时，不需要考虑从前端获取请求
		if (worker_queue.empty()) {
			zmq::poll(items.data(), 1);
		}
		else {
			zmq::poll(items);
		}

		if (items[0].revents & ZMQ_POLLIN)
		{
			/// 获取worker id
			auto replies = zmsg_recv(backend);

			/// 添加到可用worker列表中
			worker_queue.push_back(std::move(replies[0]));

			if (to_string(replies.back()) != "READY")
			{
				/// 如果是消息回复，转发到前端
				replies.erase(replies.begin(), replies.begin() + 2);
				zmsg_send(frontend, replies);
			}
		}
		else if (items[1].revents & ZMQ_POLLIN)
		{
			auto requests = zmsg_recv(frontend);

			std::vector<zmq::message_t> msgs;
			msgs.emplace_back(std::move(worker_queue.front()));
			msgs.emplace_back();
			for (auto& msg : requests)
			{
				msgs.emplace_back(std::move(msg));
			}
			zmsg_send(backend, msgs.begin(), msgs.end());

			worker_queue.pop_front();
		}
	}

	for (auto& t : threads)
	{
		t->join();
	}
	pause();
}
